#This file is part of ElectricEye.
#SPDX-License-Identifier: Apache-2.0

#Licensed to the Apache Software Foundation (ASF) under one
#or more contributor license agreements.  See the NOTICE file
#distributed with this work for additional information
#regarding copyright ownership.  The ASF licenses this file
#to you under the Apache License, Version 2.0 (the
#"License"); you may not use this file except in compliance
#with the License.  You may obtain a copy of the License at

#http://www.apache.org/licenses/LICENSE-2.0

#Unless required by applicable law or agreed to in writing,
#software distributed under the License is distributed on an
#"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#KIND, either express or implied.  See the License for the
#specific language governing permissions and limitations
#under the License.

import datetime
from check_register import CheckRegister
import base64
import json

registry = CheckRegister()

def list_tables(cache, session):
    response = cache.get("list_tables")
    if response:
        return response
    
    dynamodb = session.client("dynamodb")
    ddbTables = []
    
    for page in dynamodb.get_paginator("list_tables").paginate():
        for table in page["TableNames"]:
            ddbTables.append(
                dynamodb.describe_table(
                    TableName=table
                )
            )

    cache["list_tables"] = ddbTables
    return cache["list_tables"]

@registry.register_check("dynamodb")
def ddb_kms_cmk_check(cache: dict, session, awsAccountId: str, awsRegion: str, awsPartition: str) -> dict:
    """[DynamoDB.1] Amazon DynamoDB tables should use AWS KMS Customer Managed Keys (CMKs) for encryption at rest"""
    # ISO Time
    iso8601Time = datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat()
    for table in list_tables(cache, session):
        # B64 encode all of the details for the Asset
        assetJson = json.dumps(table,default=str).encode("utf-8")
        assetB64 = base64.b64encode(assetJson)
        tableArn = table["Table"]["TableArn"]
        tableName = table["Table"]["TableName"]
        # Attempt to assess encryption details
        try:
            sseType = table["Table"]["SSEDescription"]["SSEType"]
            if sseType == "KMS":
                cmkUsed = True
            else:
                cmkUsed = False
        except KeyError:
            cmkUsed = False
        # This is a failing check
        if cmkUsed is False:
            finding={
                "SchemaVersion": "2018-10-08",
                "Id": f"{tableArn}/ddb-kms-cmk-check",
                "ProductArn": f"arn:{awsPartition}:securityhub:{awsRegion}:{awsAccountId}:product/{awsAccountId}/default",
                "GeneratorId": f"{tableArn}/ddb-kms-cmk-check",
                "AwsAccountId": awsAccountId,
                "Types": ["Software and Configuration Checks/AWS Security Best Practices"],
                "FirstObservedAt": iso8601Time,
                "CreatedAt": iso8601Time,
                "UpdatedAt": iso8601Time,
                "Severity": {"Label": "MEDIUM"},
                "Confidence": 99,
                "Title": "[DynamoDB.1] Amazon DynamoDB tables should use AWS KMS Customer Managed Keys (CMKs) for encryption at rest",
                "Description": f"Amazon DynamoDB table {tableName} is not using an AWS KMS Customer Managed Key (CMK) for encryption. When you access an encrypted table, DynamoDB decrypts the table data transparently. You can switch between the AWS owned CMK, AWS managed CMK, and customer managed CMK at any given time. Refer to the remediation instructions if this configuration is not intended.",
                "Remediation": {
                    "Recommendation": {
                        "Text": "For more information refer to the DynamoDB Encryption at Rest section of the Amazon DynamoDB Developer Guide",
                        "Url": "https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/EncryptionAtRest.html"
                    }
                },
                "ProductFields": {
                    "ProductName": "ElectricEye",
                    "Provider": "AWS",
                    "ProviderType": "CSP",
                    "ProviderAccountId": awsAccountId,
                    "AssetRegion": awsRegion,
                    "AssetDetails": assetB64,
                    "AssetClass": "Database",
                    "AssetService": "Amazon DynamoDB",
                    "AssetComponent": "Table"
                },
                "Resources": [
                    {
                        "Type": "AwsDynamoDbTable",
                        "Id": tableArn,
                        "Partition": awsPartition,
                        "Region": awsRegion,
                        "Details": {
                            "AwsDynamoDbTable": {
                                "TableName": tableName
                            }
                        }
                    }
                ],
                "Compliance": { 
                    "Status": "FAILED",
                    "RelatedRequirements": [
                        "NIST CSF V1.1 PR.DS-1", 
                        "NIST SP 800-53 Rev. 4 MP-8",
                        "NIST SP 800-53 Rev. 4 SC-12",
                        "NIST SP 800-53 Rev. 4 SC-28",
                        "AICPA TSC CC6.1",
                        "ISO 27001:2013 A.8.2.3",
                        "CIS AWS Database Services Benchmark V1.0 4.3",
                        "CIS AWS Database Services Benchmark V1.0 4.4"
                    ]
                },
                "Workflow": {"Status": "NEW"},
                "RecordState": "ACTIVE"
            }
            yield finding
        else:
            finding={
                "SchemaVersion": "2018-10-08",
                "Id": f"{tableArn}/ddb-kms-cmk-check",
                "ProductArn": f"arn:{awsPartition}:securityhub:{awsRegion}:{awsAccountId}:product/{awsAccountId}/default",
                "GeneratorId": f"{tableArn}/ddb-kms-cmk-check",
                "AwsAccountId": awsAccountId,
                "Types": ["Software and Configuration Checks/AWS Security Best Practices"],
                "FirstObservedAt": iso8601Time,
                "CreatedAt": iso8601Time,
                "UpdatedAt": iso8601Time,
                "Severity": {"Label": "INFORMATIONAL"},
                "Confidence": 99,
                "Title": "[DynamoDB.1] Amazon DynamoDB tables should use AWS KMS Customer Managed Keys (CMKs) for encryption at rest",
                "Description": f"Amazon DynamoDB table {tableName} is using an AWS KMS Customer Managed Key (CMK) for encryption.",
                "Remediation": {
                    "Recommendation": {
                        "Text": "For more information refer to the DynamoDB Encryption at Rest section of the Amazon DynamoDB Developer Guide",
                        "Url": "https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/EncryptionAtRest.html"
                    }
                },
                "ProductFields": {
                    "ProductName": "ElectricEye",
                    "Provider": "AWS",
                    "ProviderType": "CSP",
                    "ProviderAccountId": awsAccountId,
                    "AssetRegion": awsRegion,
                    "AssetDetails": assetB64,
                    "AssetClass": "Database",
                    "AssetService": "Amazon DynamoDB",
                    "AssetComponent": "Table"
                },
                "Resources": [
                    {
                        "Type": "AwsDynamoDbTable",
                        "Id": tableArn,
                        "Partition": awsPartition,
                        "Region": awsRegion,
                        "Details": {
                            "AwsDynamoDbTable": {
                                "TableName": tableName
                            }
                        }
                    }
                ],
                "Compliance": { 
                    "Status": "PASSED",
                    "RelatedRequirements": [
                        "NIST CSF V1.1 PR.DS-1", 
                        "NIST SP 800-53 Rev. 4 MP-8",
                        "NIST SP 800-53 Rev. 4 SC-12",
                        "NIST SP 800-53 Rev. 4 SC-28",
                        "AICPA TSC CC6.1",
                        "ISO 27001:2013 A.8.2.3",
                        "CIS AWS Database Services Benchmark V1.0 4.3",
                        "CIS AWS Database Services Benchmark V1.0 4.4"
                    ]
                },
                "Workflow": {"Status": "RESOLVED"},
                "RecordState": "ARCHIVED"
            }
            yield finding

@registry.register_check("dynamodb")
def ddb_pitr_check(cache: dict, session, awsAccountId: str, awsRegion: str, awsPartition: str) -> dict:
    """[DynamoDB.2] Amazon DynamoDB tables should have Point-in-Time Recovery (PITR) enabled"""
    dynamodb = session.client("dynamodb")
    # ISO Time
    iso8601Time = datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat()
    for table in list_tables(cache, session):
        # B64 encode all of the details for the Asset
        assetJson = json.dumps(table,default=str).encode("utf-8")
        assetB64 = base64.b64encode(assetJson)
        tableArn = table["Table"]["TableArn"]
        tableName = table["Table"]["TableName"]
        # Check for PITR
        pitrCheck = dynamodb.describe_continuous_backups(
            TableName=tableName
        )["ContinuousBackupsDescription"]["PointInTimeRecoveryDescription"]["PointInTimeRecoveryStatus"]
        # this is a failing check
        if pitrCheck == "DISABLED":
            finding={
                "SchemaVersion": "2018-10-08",
                "Id": f"{tableArn}/ddb-pitr-check",
                "ProductArn": f"arn:{awsPartition}:securityhub:{awsRegion}:{awsAccountId}:product/{awsAccountId}/default",
                "GeneratorId": f"{tableArn}/ddb-pitr-check",
                "AwsAccountId": awsAccountId,
                "Types": ["Software and Configuration Checks/AWS Security Best Practices"],
                "FirstObservedAt": iso8601Time,
                "CreatedAt": iso8601Time,
                "UpdatedAt": iso8601Time,
                "Severity": {"Label": "LOW"},
                "Confidence": 99,
                "Title": "[DynamoDB.2] Amazon DynamoDB tables should have Point-in-Time Recovery (PITR) enabled",
                "Description": f"Amazon DynamoDB table {tableName} does not have Point-in-Time Recovery (PITR) enabled. Amazon DynamoDB point-in-time recovery (PITR) provides automatic backups of your DynamoDB table data, LatestRestorableDateTime is typically 5 minutes before the current time. You should consider enabling this for applications with low RTO/RPOs. Refer to the remediation instructions if this configuration is not intended.",
                "Remediation": {
                    "Recommendation": {
                        "Text": "For more information on enabling PITR refer to the Point-in-Time Recovery: How It Works section of the Amazon DynamoDB Developer Guide",
                        "Url": "https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/PointInTimeRecovery_Howitworks.html"
                    }
                },
                "ProductFields": {
                    "ProductName": "ElectricEye",
                    "Provider": "AWS",
                    "ProviderType": "CSP",
                    "ProviderAccountId": awsAccountId,
                    "AssetRegion": awsRegion,
                    "AssetDetails": assetB64,
                    "AssetClass": "Database",
                    "AssetService": "Amazon DynamoDB",
                    "AssetComponent": "Table"
                },
                "Resources": [
                    {
                        "Type": "AwsDynamoDbTable",
                        "Id": tableArn,
                        "Partition": awsPartition,
                        "Region": awsRegion,
                        "Details": {
                            "AwsDynamoDbTable": {
                                "TableName": tableName
                            }
                        }
                    }
                ],
                "Compliance": { 
                    "Status": "FAILED",
                    "RelatedRequirements": [
                        "NIST CSF V1.1 ID.BE-5",
                        "NIST CSF V1.1 PR.IP-4",
                        "NIST CSF V1.1 PR.PT-5",
                        "NIST SP 800-53 Rev. 4 CP-2",
                        "NIST SP 800-53 Rev. 4 CP-4",
                        "NIST SP 800-53 Rev. 4 CP-6",
                        "NIST SP 800-53 Rev. 4 CP-7",
                        "NIST SP 800-53 Rev. 4 CP-8",
                        "NIST SP 800-53 Rev. 4 CP-9",
                        "NIST SP 800-53 Rev. 4 CP-11",
                        "NIST SP 800-53 Rev. 4 CP-13",
                        "NIST SP 800-53 Rev. 4 PL-8",
                        "NIST SP 800-53 Rev. 4 SA-14",
                        "NIST SP 800-53 Rev. 4 SC-6",
                        "AICPA TSC A1.2",
                        "AICPA TSC A1.3",
                        "AICPA TSC CC3.1",
                        "ISO 27001:2013 A.11.1.4",
                        "ISO 27001:2013 A.12.3.1",
                        "ISO 27001:2013 A.17.1.1",
                        "ISO 27001:2013 A.17.1.2",
                        "ISO 27001:2013 A.17.1.3",
                        "ISO 27001:2013 A.17.2.1",
                        "ISO 27001:2013 A.18.1.3"
                    ]
                },
                "Workflow": {"Status": "NEW"},
                "RecordState": "ACTIVE"
            }
            yield finding
        else:
            finding={
                "SchemaVersion": "2018-10-08",
                "Id": f"{tableArn}/ddb-pitr-check",
                "ProductArn": f"arn:{awsPartition}:securityhub:{awsRegion}:{awsAccountId}:product/{awsAccountId}/default",
                "GeneratorId": f"{tableArn}/ddb-pitr-check",
                "AwsAccountId": awsAccountId,
                "Types": ["Software and Configuration Checks/AWS Security Best Practices"],
                "FirstObservedAt": iso8601Time,
                "CreatedAt": iso8601Time,
                "UpdatedAt": iso8601Time,
                "Severity": {"Label": "INFORMATIONAL"},
                "Confidence": 99,
                "Title": "[DynamoDB.2] Amazon DynamoDB tables should have Point-in-Time Recovery (PITR) enabled",
                "Description": f"Amazon DynamoDB table {tableName} does have Point-in-Time Recovery (PITR) enabled.",
                "Remediation": {
                    "Recommendation": {
                        "Text": "For more information on enabling PITR refer to the Point-in-Time Recovery: How It Works section of the Amazon DynamoDB Developer Guide",
                        "Url": "https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/PointInTimeRecovery_Howitworks.html"
                    }
                },
                "ProductFields": {
                    "ProductName": "ElectricEye",
                    "Provider": "AWS",
                    "ProviderType": "CSP",
                    "ProviderAccountId": awsAccountId,
                    "AssetRegion": awsRegion,
                    "AssetDetails": assetB64,
                    "AssetClass": "Database",
                    "AssetService": "Amazon DynamoDB",
                    "AssetComponent": "Table"
                },
                "Resources": [
                    {
                        "Type": "AwsDynamoDbTable",
                        "Id": tableArn,
                        "Partition": awsPartition,
                        "Region": awsRegion,
                        "Details": {
                            "AwsDynamoDbTable": {
                                "TableName": tableName
                            }
                        }
                    }
                ],
                "Compliance": { 
                    "Status": "PASSED",
                    "RelatedRequirements": [
                        "NIST CSF V1.1 ID.BE-5",
                        "NIST CSF V1.1 PR.IP-4",
                        "NIST CSF V1.1 PR.PT-5",
                        "NIST SP 800-53 Rev. 4 CP-2",
                        "NIST SP 800-53 Rev. 4 CP-4",
                        "NIST SP 800-53 Rev. 4 CP-6",
                        "NIST SP 800-53 Rev. 4 CP-7",
                        "NIST SP 800-53 Rev. 4 CP-8",
                        "NIST SP 800-53 Rev. 4 CP-9",
                        "NIST SP 800-53 Rev. 4 CP-11",
                        "NIST SP 800-53 Rev. 4 CP-13",
                        "NIST SP 800-53 Rev. 4 PL-8",
                        "NIST SP 800-53 Rev. 4 SA-14",
                        "NIST SP 800-53 Rev. 4 SC-6",
                        "AICPA TSC A1.2",
                        "AICPA TSC A1.3",
                        "AICPA TSC CC3.1",
                        "ISO 27001:2013 A.11.1.4",
                        "ISO 27001:2013 A.12.3.1",
                        "ISO 27001:2013 A.17.1.1",
                        "ISO 27001:2013 A.17.1.2",
                        "ISO 27001:2013 A.17.1.3",
                        "ISO 27001:2013 A.17.2.1",
                        "ISO 27001:2013 A.18.1.3"
                    ]
                },
                "Workflow": {"Status": "RESOLVED"},
                "RecordState": "ARCHIVED"
            }
            yield finding

# TODO: Deletion Protection, Autoscaling (On-Demand Capacity), Global Tables/Replicas, VPC Endpoint, CloudTrail Data Events

## END ??